#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
protocol_demo.py - 通信协议测试

Author: xiaobo (xiaobocwb@gmail.com)
Version: 1.0.0
Date: 2024-01

Copyright (c) 2024 xiaobo
Licensed under the MIT License
"""

import sys
import time
import os
import threading
import signal
import serial
import argparse
import logging
import hashlib
from tqdm import tqdm

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from protocol import Protocol

script_dir = os.path.dirname(os.path.abspath(__file__))
# firmware_path = os.path.join(script_dir, "../test/firmware.txt")
firmware_path = "E:/turbo/projects/Keil/AG@ok1054w/keil_project/Objects/AG@ok1054w.bin"


# 消息定义（需要通信双方互相约定一致）
MQTT_PUBLISH_TOPIC = 0x01
MQTT_RECEIVE_TOPIC = 0x02
APP_VERSION_TOPIC  = 0x03
OTA_INFO_TOPIC     = 0x04
OTA_DATA_TOPIC     = 0x05
OTA_END_TOPIC      = 0x06
OTA_STATE_TOPIC    = 0x07

# 缓冲区大小
BUFFER_SIZE = 512
ota_state = 0

# DTU 测试 OTA 下发
def test_dtu_ota_download(proto, version_str, file_path):
    file = open(file_path, "rb")
    file_size = os.path.getsize(file_path)
    file.seek(0, 0)
    
    # 计算文件MD5
    # ******************************************************
    md5_hash = hashlib.md5()
    for chunk in iter(lambda: file.read(4096), b""):
        md5_hash.update(chunk)
    file_md5 = md5_hash.hexdigest()
    # 将MD5值转换为字节
    file_md5_bytes = bytes.fromhex(file_md5)
    file.seek(0, 0)
    # ******************************************************

    # 发送 OTA 信息
    # ************************************************************************************************************************
    logging.info(f"DTU 发送 OTA 信息, size: {file_size} MD5: {file_md5_bytes.hex()}")
    if not proto.publish(OTA_INFO_TOPIC, file_size.to_bytes(4, "little") + file_md5_bytes + version_str.encode(), True, 0.5):
        logging.error("DTU 发送 OTA 信息失败")
        return False
    else:
        logging.info("DTU 发送 OTA 信息成功")
    # ************************************************************************************************************************

    # 等待设备 OTA 状态变为允许升级(设备可能需要等待结束当前服务)
    # **************************************************
    global ota_state
    overtime = time.time() + 60
    progress_bar = tqdm(total=60, desc='等待设备 OTA 状态变为 1')
    while time.time() < overtime and 0 == ota_state:
        proto.publish(OTA_STATE_TOPIC, b"", False)
        time.sleep(1)
        progress_bar.update(1)
        progress_bar.set_postfix(state=ota_state)

    progress_bar.close()

    if ota_state != 1:
        return
    # **************************************************

    # 发送OTA固件数据
    # *******************************************************************************************
    send_size = 0
    progress_bar = tqdm(total=file_size, unit='B', unit_scale=True, desc='DTU 发送进度')
    while True:
        data = file.read(BUFFER_SIZE)
        if not data:
            break

        if not proto.publish(OTA_DATA_TOPIC, data, True, 15):
            logging.error("DTU 发送失败")
            return False

        send_size += len(data)
        progress_bar.update(len(data))
        progress_bar.set_postfix(send_size=send_size)

        time.sleep(0.001)

    progress_bar.close()
    # *******************************************************************************************

    time.sleep(3)

    # 发送OTA结束信息
    # *******************************************************************************************
    if not proto.publish_try(OTA_END_TOPIC, b"", True, 200, 3):
        logging.error("DTU 发送结束失败")
        return False
    else:
        logging.info("DTU 发送结束成功")
    # *******************************************************************************************

    file.close()
    return True


# 接收线程
def recv_thread_func(proto, serial):
    while True:
        data = serial.read(BUFFER_SIZE)
        if data:
            proto.parse(data)


if __name__ == "__main__":
    # parser = argparse.ArgumentParser(description='UART Protocol')
    # parser.add_argument('port', type=str, help='UART port', default="COM11")
    # parser.add_argument('baudrate', type=int, help='UART baudrate', default=115200)

    # args = parser.parse_args()
    # port = args.port
    # baudrate = args.baudrate

    port = 'COM11'
    baudrate = 115200

    # 初始化串口和协议
    serial = serial.Serial(port, baudrate, parity=serial.PARITY_NONE, stopbits=serial.STOPBITS_ONE, timeout=0.001)
    proto = Protocol(serial.write)

    # 启动接收线程
    thread = threading.Thread(target=recv_thread_func, args=(proto, serial))
    thread.start()

    # 订阅消息
    proto.subscribe(MQTT_PUBLISH_TOPIC, lambda topic, payload: logging.info(f"DTU 收到MQTT消息: msgid={topic}, len={len(payload)}, payload={int.from_bytes(payload, 'little')}"))
    proto.subscribe(MQTT_RECEIVE_TOPIC, lambda topic, payload: logging.info(f"DTU 收到MQTT消息: msgid={topic}, len={len(payload)}, payload={payload.hex()}"))
    proto.subscribe(APP_VERSION_TOPIC, lambda topic, payload: logging.info(f"DTU 收到APP版本: msgid={topic}, len={len(payload)}, payload={payload.decode('utf-8')}"))
    proto.subscribe(OTA_STATE_TOPIC, lambda topic, payload: handle_ota_state(topic, payload))

    def handle_ota_state(topic, payload):
        global ota_state
        ota_state = int.from_bytes(payload, 'little')

    proto.publish(APP_VERSION_TOPIC, b"", True, 0.5)

    # 测试代码
    if not test_dtu_ota_download(proto, 'v1.1.2', firmware_path):
        logging.error("DTU OTA 测试失败")
    else:
        logging.info("DTU OTA 测试成功")
